"""
分割分析服务

这是jinxiang_FF模块的主要接口，封装SegmentationAnalyzer，提供简化的文件夹分析功能。
"""

import os
import sys
import time
import numpy as np
from typing import Dict, Any, List, Optional
from pathlib import Path
from PIL import Image
import matplotlib.pyplot as plt
import cv2
# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
sys.path.insert(0, project_root)

from devdeploy.src import SegmentationAnalyzer, get_optimal_device, prepare_output_directory

from .jx_ff_expert_alg import visualize_complete_analysis, filter_contours_by_area_and_aspect_ratio, sliding_window_analysis, analyze_final_level

class SegmentationAnalyzerService:
    """
    分割分析服务
    
    这是jinxiang_ZZ模块的主要接口，提供简化的文件夹分析功能。
    """
    
    def __init__(self, model_path: str, device: Optional[str] = None):
        """
        初始化分割分析服务
        
        Args:
            model_path (str): 模型文件路径
            device (str, optional): 推理设备，如果为None则自动选择
        """
        self.model_path = model_path
        self.device = device or get_optimal_device()
        
        # 初始化分割分析器
        try:
            self.analyzer = SegmentationAnalyzer(
                model_path=self.model_path,
                device=self.device
            )
            print(f"分割分析服务初始化成功，使用设备: {self.device}")
        except Exception as e:
            raise RuntimeError(f"初始化分割分析服务失败: {str(e)}")
    
    def analyze_oneimg(self, image_path: str, output_folder: str, target_classes: List[int] = [1]) -> Dict[str, Any]:
        """
        分析文件夹中的所有图像
        
        Args:
            input_folder (str): 输入文件夹路径
            output_folder (str): 输出文件夹路径
            target_classes (List[int]): 要分析的类别，默认为[1, 2]
            
        Returns:
            Dict[str, Any]: 分析结果
        """

        result = {
            'success': False,
            'message': '',
            'result': 1.5,
           
        }
        comparison=None
        paint_result=None
        try:

            # 分析单张图像（计时）
            start_time = time.time()
            analysis_result = self.analyzer.analyze_single(image_path, target_classes)

            '''
                    analysis_result = {
            'class_ratios': analysis_result['class_ratios'],
            'class_masks': analysis_result['class_masks'],
            'class_names': class_names,
            'image_shape': result['segmentation_mask'].shape,
            'segmentation_mask': result['segmentation_mask']
            }
            '''
            elapsed_analyze = time.time() - start_time
            print(f"  - 单图像分析耗时: {elapsed_analyze:.3f} 秒")
            
            # opencv  根据轮廓的面积大小 ， 长宽比大小  对分割二值图matimg进行初步过滤 ，并返回筛选后的二值图
            matimg = analysis_result['segmentation_mask']
                  
            #matimg  从int64转化为int8
            matimg=matimg.astype(np.uint8)

            final_level=0.5

            filtered_img, filtered_contours=filter_contours_by_area_and_aspect_ratio(matimg)

            # 根据过滤后的二值图像filtered_img， 以h*w的滑窗形式  计算每个滑窗中图片根据 _result_quantization 函数计算带状组织等级，并记录每个滑窗的等级，最后综合分析每个滑窗的等级，得到最终的等级
            window_results, level_map=sliding_window_analysis(filtered_img)
            

            #综合分析带状组织等级
            final_level, level_stats=analyze_final_level(window_results, 'worst_case')
            
            #将levelmap  绘制到原图上
            original_img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
            
            overlay_img, comparison,paint_result=visualize_complete_analysis(original_img, filtered_img, window_results, level_map, final_level, level_stats)
            # 获取文件名（不含扩展名）
            file_name = Path(image_path).stem
            
          
            result = {
                'success': True,
                'message': '',
                'result': final_level,
            
            }
            
        except Exception as e:
            print(f"处理图像失败 {image_path}: {str(e)}")
            result = {
            'success': False,
            'message': f"处理图像失败 {image_path}: {str(e)}",
            'result':0.5,
           
        }
            return result,comparison,paint_result 
        return result,comparison,paint_result

    def analyze_folder(self, input_folder: str, output_folder: str, 
                      target_classes: List[int] = [1, 2]) -> Dict[str, Any]:
        """
        分析文件夹中的所有图像
        
        Args:
            input_folder (str): 输入文件夹路径
            output_folder (str): 输出文件夹路径
            target_classes (List[int]): 要分析的类别，默认为[1, 2]
            
        Returns:
            Dict[str, Any]: 分析结果
        """
        result = {
            'success': False,
            'message': '',
            'total_images': 0,
            'processed_images': 0,
            'failed_images': 0
        }
        
        try:
            # 验证输入文件夹
            if not os.path.exists(input_folder):
                result['message'] = f"输入文件夹不存在: {input_folder}"
                return result
            
            if not os.path.isdir(input_folder):
                result['message'] = f"路径不是文件夹: {input_folder}"
                return result
            
            # 使用prepare_output_directory确保输出文件夹正确创建
            if not prepare_output_directory(output_folder, clear_existing=True):
                result['message'] = f"无法创建输出文件夹: {output_folder}"
                return result
            
            # 创建result子文件夹用于保存图像结果
            result_folder = output_folder
            
            # 获取图像文件列表
            image_files = self._get_image_files(input_folder)
            result['total_images'] = len(image_files)
            
            if len(image_files) == 0:
                result['message'] = f"文件夹中没有找到支持的图像文件: {input_folder}"
                return result
            
            print(f"开始分析 {len(image_files)} 个图像文件...")
            
            # 存储所有分析结果用于生成报告
            all_results = []

            # 逐个处理图像
            for i, image_path in enumerate(image_files):
                try:
                    print(f"处理图像 {i+1}/{len(image_files)}: {os.path.basename(image_path)}")
                    
                    # 分析单张图像（计时）
                    start_time = time.time()
                    analysis_result = self.analyzer.analyze_single(image_path, target_classes)

                    '''
                            analysis_result = {
                    'class_ratios': analysis_result['class_ratios'],
                    'class_masks': analysis_result['class_masks'],
                    'class_names': class_names,
                    'image_shape': result['segmentation_mask'].shape,
                    'segmentation_mask': result['segmentation_mask']
                    }
                    '''
                    elapsed_analyze = time.time() - start_time
                    print(f"  - 单图像分析耗时: {elapsed_analyze:.3f} 秒")

                    # 获取文件名（不含扩展名）
                    file_name = Path(image_path).stem
                    
                    count=self._result_quantization(analysis_result)
                    # 创建可视化结果（计时）
                    vis_start_time = time.time()
                    visualization_path = os.path.join(result_folder, f"{file_name}_visualization.png")
                    self._create_visualization(image_path, analysis_result, visualization_path)
                    elapsed_vis = time.time() - vis_start_time
                    print(f"  - 可视化生成耗时: {elapsed_vis:.3f} 秒")
                    
                    # 记录结果用于报告，包含类别名称信息（计时）
                    record_start_time = time.time()
                    all_results.append({
                        'file_name': file_name,
                        'class_ratios': analysis_result['class_ratios'],
                        'class_names': analysis_result.get('class_names', {}),
                        'image_shape': analysis_result['image_shape']
                    })
                    elapsed_record = time.time() - record_start_time
                    print(f"  - 结果记录耗时: {elapsed_record:.3f} 秒")
                    
                    result['processed_images'] += 1
                    
                except Exception as e:
                    print(f"处理图像失败 {image_path}: {str(e)}")
                    result['failed_images'] += 1
            
            # 生成汇总报告到根目录
            self._generate_summary_report(all_results, output_folder, target_classes)
            
            result['success'] = True
            result['message'] = f"分析完成！成功处理 {result['processed_images']} 个文件，失败 {result['failed_images']} 个文件"
            
        except Exception as e:
            result['message'] = f"分析文件夹时发生错误: {str(e)}"
        
        return result
    
    def _get_image_files(self, folder_path: str) -> List[str]:
        """
        获取文件夹中的所有图像文件
        
        Args:
            folder_path (str): 文件夹路径
            
        Returns:
            List[str]: 图像文件路径列表
        """
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
        image_files = []
        
        for file_path in Path(folder_path).iterdir():
            if file_path.is_file() and file_path.suffix.lower() in image_extensions:
                image_files.append(str(file_path))
        
        return sorted(image_files)
    
    
            
    def _generate_summary_report(self, all_results: List[Dict], output_folder: str, target_classes: List[int]):
        """
        生成汇总报告到根目录的result.txt文件
        
        Args:
            all_results (List[Dict]): 所有分析结果
            output_folder (str): 输出文件夹
            target_classes (List[int]): 目标类别
        """
        try:
            summary_path = os.path.join(output_folder, 'result.txt')
            
            with open(summary_path, 'w', encoding='utf-8') as f:
                f.write("金相组织分析汇总报告\n")
                f.write("=" * 50 + "\n")
                f.write(f"分析时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"模型路径: {self.model_path}\n")
                f.write(f"目标类别: {target_classes}\n")
                f.write(f"总文件数: {len(all_results)}\n\n")
                
                if all_results:
                    # 获取类别名称映射（从第一个结果中获取）
                    class_names = all_results[0].get('class_names', {})
                    
                    # 计算平均占比
                    avg_ratios = {}
                    for class_id in target_classes:
                        ratios = [r['class_ratios'].get(class_id, 0) for r in all_results]
                        avg_ratios[class_id] = sum(ratios) / len(ratios)
                    
                    f.write("整体统计:\n")
                    for class_id, avg_ratio in avg_ratios.items():
                        class_name = class_names.get(class_id, f"类别{class_id}")
                        f.write(f"{class_name} (ID:{class_id}) 平均占比: {avg_ratio:.2f}%\n")
                    f.write("\n")
                    
                    f.write("详细结果:\n")
                    for r in all_results:
                        f.write(f"{r['file_name']} - ")
                        for class_id in target_classes:
                            ratio = r['class_ratios'].get(class_id, 0)
                            class_name = class_names.get(class_id, f"类别{class_id}")
                            f.write(f"{class_name}: {ratio:.2f}% ")
                        f.write("\n")
            
            print(f"汇总报告已保存: {summary_path}")
            
        except Exception as e:
            print(f"生成汇总报告失败: {str(e)}")
    
    def _create_visualization(self, original_image_path: str, result: Dict[str, Any], output_path: str) -> bool:
        """
        创建简化的可视化结果，只保存overlap图像以提高性能
        
        Args:
            original_image_path: 原始图像路径
            result: 分析结果
            output_path: 输出路径
            
        Returns:
            bool: 是否保存成功
        """
        try:
            # 加载原始图像
            img = Image.open(original_image_path)
            
            # 获取结果数据
            class_ratios = result['class_ratios']
            class_masks = result['class_masks']
            
            # 创建overlap图像（P类和F类叠加）
            overlap_img = self._create_overlap_visualization(img, class_masks, class_ratios)
            
            if overlap_img is None:
                print(f"无法创建overlap图像: {output_path}")
                return False
            
            # 创建单个子图显示overlap图像
            fig, ax = plt.subplots(1, 1, figsize=(10, 8))
            
            # 显示overlap图像
            ax.imshow(overlap_img)
            ax.set_title('P+F Overlap')
            ax.axis('off')
            
            # 添加统计信息
            stats_text = f"P: {class_ratios.get(1, 0):.2f}%\nF: {class_ratios.get(2, 0):.2f}%"
            ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, 
                   fontsize=14, fontweight='bold', color='white',
                   bbox=dict(boxstyle="round,pad=0.3", facecolor="black", alpha=0.7),
                   verticalalignment='top')
            
            # 保存图像
            plt.tight_layout()
            plt.savefig(output_path, dpi=150, bbox_inches='tight')
            plt.close()
            
            return True
            
        except Exception as e:
            print(f"创建可视化失败 {output_path}: {str(e)}")
            return False
    
    def _create_overlap_visualization(self, original_img: Image.Image, class_masks: Dict, class_ratios: Dict) -> Optional[np.ndarray]:
        """
        创建P类和F类的overlap可视化
        
        Args:
            original_img: 原始图像
            class_masks: 类别掩码字典
            class_ratios: 类别占比字典
            
        Returns:
            np.ndarray: overlap图像
        """
        try:
            # 将PIL图像转换为numpy数组
            if isinstance(original_img, Image.Image):
                img_array = np.array(original_img)
            else:
                img_array = original_img
            
            # 创建overlap图像
            overlap_img = img_array.copy()
            
            # 获取P类和F类的掩码
            if 1 not in class_masks or 2 not in class_masks:
                return None
                
            p_mask = class_masks[1]
            f_mask = class_masks[2]
            
            # 创建彩色overlay
            # P类用红色，F类用绿色，重叠区域用黄色
            p_overlay = np.zeros_like(img_array)
            f_overlay = np.zeros_like(img_array)
            
            # P类区域（红色）
            p_overlay[p_mask > 0] = [255, 0, 0]  # 红色
            # F类区域（绿色）
            f_overlay[f_mask > 0] = [0, 255, 0]  # 绿色
            
            # 重叠区域（黄色）
            overlap_mask = (p_mask > 0) & (f_mask > 0)
            p_overlay[overlap_mask] = [255, 255, 0]  # 黄色
            f_overlay[overlap_mask] = [255, 255, 0]  # 黄色
            
            # 将overlay与原图混合
            alpha = 0.6  # 透明度
            overlap_img = (1 - alpha) * img_array + alpha * (p_overlay + f_overlay)
            overlap_img = np.clip(overlap_img, 0, 255).astype(np.uint8)
            
            return overlap_img
            
        except Exception as e:
            print(f"Create overlap visualization failed: {str(e)}")
            return None